"""
WebSearchResearchAgent - 主模块
基于PocketFlow框架的搜索研究代理
"""

from pocketflow import Flow
from nodes import WebSearchNode, DecideNode, AnswerQuestionNode
from config import get_config, update_config
from logger import AgentLogger

def create_flow():
    """
    创建代理流程
    
    返回:
        配置好的Flow实例
    """
    # 获取配置
    config = get_config()
    
    # 创建节点实例
    search_node = WebSearchNode(max_retries=2)
    decide_node = DecideNode(max_retries=2)
    answer_node = AnswerQuestionNode(max_retries=2)
    
    # 配置流程
    search_node - "decide" >> decide_node  # 明确指定"decide"而不是默认的"default"
    decide_node - "search" >> search_node
    decide_node - "answer" >> answer_node
    
    # 创建流程实例
    flow = Flow(start=decide_node)
    
    return flow

def run_agent(question=None, search_query=None, debug=False):
    """
    运行搜索研究代理
    
    参数:
        question: 要回答的问题，如果为None则会提示输入
        search_query: 初始搜索关键词，如果为None则会使用问题或提示输入
        debug: 是否开启调试模式
    
    返回:
        包含代理运行结果的shared字典
    """
    # 获取配置
    config = get_config()
    
    # 更新调试模式
    if debug:
        update_config("flow", "debug_mode", True)
        update_config("log", "level", "DEBUG")
    
    # 创建流程
    flow = create_flow()
    
    # 初始化共享存储
    shared = {}
    
    # 创建日志记录器
    logger = AgentLogger(config["log"])
    shared["logger"] = logger
    
    # 设置问题和搜索查询
    if question:
        shared["question"] = question
        logger.info(f"设置问题: {question}")
    
    if search_query:
        shared["search_query"] = search_query
        logger.info(f"设置初始搜索查询: {search_query}")
    elif question:
        shared["search_query"] = question
        logger.info(f"使用问题作为初始搜索查询: {question}")
    
    # 运行流程
    logger.info("启动搜索研究代理")
    print("🚀 启动搜索研究代理")
    
    try:
        flow.run(shared)
        logger.info("流程正常结束")
    except Exception as e:
        logger.error(f"流程运行出错: {str(e)}")
        import traceback
        logger.error(traceback.format_exc())
        print(f"❌ 运行出错: {str(e)}")
    finally:
        # 保存历史记录
        logger.save()
    
    print("🏁 流程结束")
    return shared

if __name__ == "__main__":
    import argparse
    
    # 解析命令行参数
    parser = argparse.ArgumentParser(description="Web搜索研究代理")
    parser.add_argument("-q", "--question", help="要回答的问题")
    parser.add_argument("-s", "--search", help="初始搜索关键词")
    parser.add_argument("-d", "--debug", action="store_true", help="开启调试模式")
    args = parser.parse_args()
    
    # 运行代理
    result = run_agent(
        question=args.question,
        search_query=args.search,
        debug=args.debug
    )
    
    # 如果想要保存最终答案到文件
    if "answer" in result:
        with open("answer.txt", "w", encoding="utf-8") as f:
            f.write(result["answer"])
        print("✅ 答案已保存到 answer.txt")


